import torch
import torch.nn as nn
import spconv
from torch.nn.modules.upsampling import Upsample

device = torch.device('cuda')

# U-net
# 
class conv_block(nn.Module):
    def __init__(self,ch_in,ch_out):
        super(conv_block,self).__init__()

        self.conv = spconv.SparseSequential(
            spconv.SparseConv3d(ch_in, ch_out, kernel_size=3,stride=1,padding=1),
            # nn.BatchNorm3d(ch_out),
            nn.ReLU(inplace=True),
            spconv.SparseConv3d(ch_out, ch_out, kernel_size=3,stride=1,padding=1),
            # nn.BatchNorm3d(ch_out),
            nn.ReLU(inplace=True),
            # spconv.ToDense()
        )

        # self.conv = nn.Sequential(
        #     nn.Conv3d(ch_in, ch_out, kernel_size=3,stride=1,padding=1,bias=True),
        #     nn.BatchNorm3d(ch_out),
        #     nn.ReLU(inplace=True),
        #     nn.Conv3d(ch_out, ch_out, kernel_size=3,stride=1,padding=1,bias=True),
        #     nn.BatchNorm3d(ch_out),
        #     nn.ReLU(inplace=True)
        # )


    def forward(self, x):
        # x = x.to_sparse(x.ndim - 1)
        # spatial_shape = x.shape[1:-1]
        # batch_size = x.shape[0]
        # indices_th = x.indices().permute(1, 0).contiguous().int()
        # features_th = x.values()
        # #return cls(features_th, indices_th, spatial_shape, batch_size)
        # x_sp = spconv.SparseConvTensor(features_th, indices_th, spatial_shape, batch_size)
        # x_sp = spconv.SparseConvTensor.from_dense(x)
        # print('conv1',x.shape)
        output = self.conv(x)
        return output
        # x_sp = spconv.SparseConvTensor.from_dense(x)
        # output = self.conv(x_sp)
        # return output

class inv_conv_block(nn.Module):
    def __init__(self,ch_in,ch_out):
        super(inv_conv_block,self).__init__()

        self.inv_conv = spconv.SparseSequential(
            spconv.SparseInverseConv3d(ch_in,ch_out,kernel_size=3,indice_key='cp0'),
            nn.BatchNorm3d(ch_out),
            nn.ReLU(inplace=True),
            spconv.SparseInverseConv3d(ch_in,ch_out,kernel_size=3,indice_key='cp0'),
            nn.BatchNorm3d(ch_out),
            nn.ReLU(inplace=True),
            # spconv.ToDense(),
        )

        # self.conv = nn.Sequential(
        #     nn.Conv3d(ch_in, ch_out, kernel_size=3,stride=1,padding=1,bias=True),
        #     nn.BatchNorm3d(ch_out),
        #     nn.ReLU(inplace=True),
        #     nn.Conv3d(ch_out, ch_out, kernel_size=3,stride=1,padding=1,bias=True),
        #     nn.BatchNorm3d(ch_out),
        #     nn.ReLU(inplace=True)
        # )


    def forward(self, x: torch.Tensor):
        # x = x.to_sparse(x.ndim - 1)
        # spatial_shape = x.shape[1:-1]
        # batch_size = x.shape[0]
        # indices_th = x.indices().permute(1, 0).contiguous().int()
        # features_th = x.values()
        # #return cls(features_th, indices_th, spatial_shape, batch_size)
        # x_sp = spconv.SparseConvTensor(features_th, indices_th, spatial_shape, batch_size)
        # x_sp = spconv.SparseConvTensor.from_dense(x)
        output = self.inv_conv(x)
        return output


class up_conv(nn.Module):
    def __init__(self,ch_in,ch_out):
        super(up_conv,self).__init__()

        self.up = spconv.SparseSequential(
            nn.Upsample(scale_factor=2),
            spconv.SparseInverseConv3d(ch_in,ch_out,kernel_size=3,indice_key='cp0'),
            nn.BatchNorm3d(ch_out),
            nn.ReLU(inplace=True),
            # spconv.ToDense()
            
        )
        # self.up = nn.Sequential(
        #     nn.Upsample(scale_factor=2),
        #     nn.Conv3d(ch_in,ch_out,kernel_size=3,stride=1,padding=1,bias=True),
		#     nn.BatchNorm3d(ch_out),
		# 	  nn.ReLU(inplace=True)
        # )

    def forward(self, x: torch.Tensor):
        # x = x.to_sparse(x.ndim - 1)
        # spatial_shape = x.shape[1:-1]
        # batch_size = x.shape[0]
        # indices_th = x.indices().permute(1, 0).contiguous().int()
        # features_th = x.values()
        # #return cls(features_th, indices_th, spatial_shape, batch_size)
        # x_sp = spconv.SparseConvTensor(features_th, indices_th, spatial_shape, batch_size)
        # x_sp = spconv.SparseConvTensor.from_dense(x)
        output = self.up(x)
        return output
        # x_sp = spconv.SparseConvTensor.from_dense(x)
        # output = self.up(x_sp)
        # return output


class SparseMaxPool(nn.Module):
    def __init__(self, kernel_s, stride):
        super(SparseMaxPool, self).__init__()

        self.maxp = spconv.SparseSequential(
            spconv.SparseMaxPool3d(kernel_size=kernel_s),
            # spconv.ToDense()
        )

    def forward(self, x: torch.Tensor):
        # x = x.to_sparse(x.ndim - 1)
        # spatial_shape = x.shape[1:-1]
        # batch_size = x.shape[0]
        # indices_th = x.indices().permute(1, 0).contiguous().int()
        # features_th = x.values()
        # #return cls(features_th, indices_th, spatial_shape, batch_size)
        # x_sp = spconv.SparseConvTensor(features_th, indices_th, spatial_shape, batch_size)
        # x_sp = spconv.SparseConvTensor.from_dense(x)
        output = self.maxp(x)
        return output
        # x_sp = spconv.SparseConvTensor.from_dense(x)
        # output = self.maxp(x_sp)
        # return output



class U_Net(nn.Module):
    def __init__(self,init_ch=36,output_ch=18):
        super(U_Net,self).__init__()
        
        self.Maxpool = SparseMaxPool(kernel_s=2,stride=2)

        # 36 -> 18
        self.Conv1 = conv_block(ch_in=init_ch,ch_out=64)
        # 18 -> 9
        self.Conv2 = conv_block(ch_in=64,ch_out=128)
        # # 9 -> 5
        self.Conv3 = conv_block(ch_in=128,ch_out=256)
        # self.Conv4 = conv_block(ch_in=256,ch_out=512)
        # self.Conv5 = conv_block(ch_in=512,ch_out=1024)

        # self.Up5 = up_conv(ch_in=1024,ch_out=512)
        # self.Up_conv5 = conv_block(ch_in=1024, ch_out=512)

        # self.Up4 = up_conv(ch_in=512,ch_out=256)
        # self.Up_conv4 = conv_block(ch_in=512, ch_out=256)
        
        self.Up3 = up_conv(ch_in=256,ch_out=128)
        self.Up_conv3 = inv_conv_block(ch_in=256, ch_out=128)
        
        self.Up2 = up_conv(ch_in=128,ch_out=64)
        self.Up_conv2 = inv_conv_block(ch_in=128, ch_out=64)

        self.Conv_1x1 = nn.Conv3d(64,output_ch,kernel_size=1,stride=1,padding=0)


    def forward(self,x:torch.Tensor):
        print(x.shape)
        # encoding path
        x = spconv.SparseConvTensor.from_dense(x)
        print('sparse_shape', x.spatial_shape)
       

        x1 = self.Conv1(x)

        x2 = self.Maxpool(x1)
        x2 = self.Conv2(x2)
        
        x3 = self.Maxpool(x2)
        x3 = self.Conv3(x3)

        d3 = self.Up3(x3)
        d3 = spconv.JoinTable(x2,d3)
        # # jointable
        # d3 = torch.cat((x2,d3),dim=1)
        d3 = self.Up_conv3(d3)

        d2 = self.Up2(d3)

        d2 = spconv.JoinTable(x1,d2)
        # # jointable
        # d2 = torch.cat((x1,d2),dim=1)
        d2 = self.Up_conv2(d2)

        d1 = self.Conv_1x1(d2)

        return d1


